import gradio as gr
import numpy as np
from audio.audio_processor import AudioProcessor
from debate_controller import DebateController
import time
from audio.chattts import text_to_speech
from pydub import AudioSegment
import threading
import queue
from audio.asr_model import asr_decode, splice_chunks

# 在文件顶部添加CSS定义
css = """
.vertical-text{
    height: 550px !important;
    display: flex !important;
    flex-direction: column !important;
    gap: 20px !important;  /* 控制元素间距 */
}

/* 为包含HighlightedText的容器添加样式 */
.speakers-container {
    display: flex !important;    height: 100% !important;

    flex-direction: column !important;
    justify-content: flex-end !important;
}
"""
js = """
                () => {
                    setInterval(() => {
                        const chatbot = document.querySelector('.gradio-chatbot');
                        if(chatbot) {
                            chatbot.scrollTop = chatbot.scrollHeight;
                        }
                    }, 100);
                }
                """


def main():
    with gr.Blocks(title="智能辩论系统", css=css, js=js) as demo:

        # region === 界面元素 ===
        # 在main函数开始处添加状态管理
        debate_state = gr.State({
            'running': False,
            'speaker': "主持人",
            'debate_state': [],  # 新增音频队列
            'audio_count': 0,  # 新增音频片段计数器
            'play_voice': False,
            'debate_model': "ai_vs_ai",
            'human_size': "正方",
            'topic': '',
            'current_human_debater': '',
            'current_human_debater_message': '',
            'waiting_input': False,
            'current_role': "assistant",  # 初始角色
            'recording': False,
            'record_cache': np.array([]),
            'agent_type': 'basic'
        })

        # === 标题行 ===
        gr.Markdown("<h1 style='text-align: center;'>基于多智能体的模拟辩论系统</h1>")

        # === 配置行 ===
        with gr.Row():
            debate_mode = gr.Dropdown(
                label="辩论模式",
                choices=["智能体辩论", "人机辩论"],
                value="智能体辩论"
            )
            debate_topic = gr.Textbox(
                label="辩论题目",
                placeholder="请输入辩论主题...",
                value="手机应该进入中小学校园",
            )
            agent_type = gr.Dropdown(
                label="辩论智能体",
                choices=["简单智能体", "RAG反思智能体"],
                value="简单智能体"
            )

        # === 控制按钮 ===
        with gr.Row():
            start_btn = gr.Button("▶️ 开始")
            stop_btn = gr.Button("⏹️ 停止", interactive=False)
            reset_btn = gr.Button("🔄 重启")

        # === 主内容区 ===
        with gr.Row():
            # 左侧辩手面板
            with gr.Column(scale=1):
                pro_panel_label = gr.Markdown("### 正方")  # 添加变量引用
                pro_role = gr.Dropdown(
                    label="角色选择",
                    choices=["正方", "反方"],
                    value="正方",
                    interactive=False
                )
                with gr.Column(elem_classes="speakers-container"):
                    pro_speakers = gr.HighlightedText(
                        label="当前辩手",
                        elem_classes=["vertical-text"],
                        value=[("正方一辩", "蓄势待发"), ("正方二辩", "蓄势待发")],
                        color_map={"正在发言": "green", "蓄势待发": "gray"},
                        combine_adjacent=True,
                        adjacent_separator='\n'
                    )

                # 中央辩论大屏
            with gr.Column(scale=3):
                debate_screen = gr.Chatbot(
                    label="辩论实况",
                    height=630,
                    type='messages',
                    render_markdown=True,  # 启用Markdown渲染
                    bubble_full_width=False
                )
                is_play = gr.Checkbox(label="播放语音", value=False)
                soundfile = gr.Audio(
                    label="语音输出",
                    visible=False,
                    autoplay=True
                )
            # 右侧动态面板
            with gr.Column(scale=1, elem_classes="right-panel"):
                right_panel_label = gr.Markdown("### 反方")
                con_role = gr.Dropdown(
                    label="角色选择",
                    choices=["反方", "正方"],
                    value="反方",
                    interactive=False
                )
                with gr.Column():
                    con_speakers = gr.HighlightedText(
                        label="当前辩手",
                        elem_classes=["vertical-text"],
                        value=[("反方一辩", "蓄势待发"), ("反方二辩", "蓄势待发")],
                        color_map={"正在发言": "red", "蓄势待发": "gray"},
                        combine_adjacent=True,
                        adjacent_separator='\n'

                    )

                # 人机交互组件（默认隐藏）
                with gr.Column(visible=False) as human_interface:
                    human_input = gr.Textbox(
                        label="人类辩手输入",
                        placeholder="当上方“当前辩手”高亮时，请对应的辩手录音发言或者手动输入，最后按发送键提交辩词...",
                        lines=12,
                    )
                    with gr.Row():
                        record_audio = gr.Audio(sources=["microphone"], streaming=True)  # 修改后的Audio组件
                        send_btn = gr.Button("发送", interactive=False)

        # endregion

        # region === 交互逻辑 ===
        def update_role(role):
            opposite = "反方" if role == "正方" else "正方"
            return [
                gr.update(value=opposite),  # 更新对方下拉框
                gr.update(value=f"### {role}"),  # 更新本方标题
                gr.update(value=f"### {opposite}")  # 更新对方标题
            ]

        def update_agent_list(agent):
            return gr.update(value=f"### {'人类' if agent == '人机辩论' else '反方'}")  # 修改这行

        def update_right_panel(mode, state):
            if mode == "人机辩论":
                state['debate_model'] = "human_vs_ai"
                return [
                    gr.update(visible=True),  # 显示语言输入等组件
                    gr.update(elem_classes=""),  # 缩小高亮小面板
                    gr.update(interactive=True),
                    gr.update(interactive=True),
                    state,
                ]
            else:
                state['debate_model'] = "ai_vs_ai"
                return [
                    gr.update(visible=False),  # 更新对方下拉框
                    gr.update(elem_classes=["vertical-text"]),
                    gr.update(interactive=False),
                    gr.update(interactive=False),
                    state,
                ]

        # endregion

        # region 状态定义
        def reset_debate(state):
            state['running'] = False
            state['controller'] = None
            return [
                [],  # 对应debate_screen
                state,
                gr.update(interactive=True),  # start_btn
                gr.update(interactive=False),  # stop_btn
                gr.update(value="", interactive=True),  # debate_topic
                gr.update(interactive=True),  # debate_mode
                gr.update(interactive=True),  # agent_type
                gr.update(interactive=True),  # pro_role
                gr.update(interactive=True),  # con_role
                gr.update(interactive=True),  # is_play
            ]


        # endregion
        def update_interactive(state, speaker):
            for size in state['debate_state']:
                for debater_info in size:
                    if debater_info[0] == speaker:
                        debater_info[1] = "正在发言"
                    else:
                        debater_info[1] = "蓄势待发"

        def show_text(history, state, current_role, message):
            history.append({"role": current_role, "content": ""})
            for character in message:
                if state["running"] is False:
                    break
                history[-1]['content'] += character
                time.sleep(0.05)
                # yield history
                yield history, state, None, state['debate_state'][0], state['debate_state'][1], gr.update(interactive=False)

        def show_audio(history, state, current_role, message):
            audio_processor = AudioProcessor(message)
            state['audio_count'] += audio_processor.audio_count  # 记录总音频段数
            # 立即更新聊天记录
            while state['audio_count'] > 0:
                if state['running'] is False:
                    break
                audio_info, sleep_time = audio_processor.get_audio()
                if audio_info:
                    print(f"audio_info: {audio_info[0]}")
                    history.append({"role": current_role, "content": audio_info[2]})
                    yield history, state, audio_info[1], state['debate_state'][0], state['debate_state'][1], gr.update(interactive=False)
                    state['audio_count'] -= 1
                    time.sleep(sleep_time)

        def debate_generator(topic, history: list, state: dict):
            topic = topic.strip()
            if topic == "":
                return [], state
            state['topic'] = topic
            # 从state中恢复控制器或新建
            if 'controller' not in state or state['controller'] is None:  # 新增状态保存点
                controller = DebateController(topic,
                                              debate_model=state['debate_model'],
                                              human_size=state['human_size'], agent_type=state['agent_type'])
                state['controller'] = controller
                state["debate_state"] = controller.get_debates_info()
                # 将生成器实例保存到state中
                debate_gen = controller.run_debate()
                state['current_generator'] = debate_gen  # 保存生成器实例
                # 当人机辩论时，辩手面板显示随着人类辩手的正反方切换
                if  state['debate_model'] == 'human_vs_ai' and state['human_size'] == "正方":
                    state['debate_state'][0], state['debate_state'][1] = state['debate_state'][1], \
                    state['debate_state'][0]
            else:
                controller = state['controller']  # 复用已有控制器



            try:
                if state['waiting_input']:
                   controller.human_response(state['current_human_debater_message'])
                   state['waiting_input'] = False
                   state['record_cache'] = np.array([])
                # 直接使用控制器的生成器

                while not state['waiting_input']:
                    speaker, message =next(state['current_generator'])  # 新增恢复参数
                    if state['running'] is False:
                        break
                    state['current_role'] = "user" if state['current_role'] == "assistant" else "assistant"
                    update_interactive(state, speaker)

                    if message == "human_turn":
                        state['current_human_debater'] = speaker
                        state['waiting_input'] = True  # 设置恢复标志
                        yield history, state, None, state['debate_state'][0], state['debate_state'][1], gr.update(interactive=True)
                        return  # 正常退出生成器

                    playing = state['play_voice']
                    if playing:
                        yield from show_audio(history, state, state['current_role'], message)
                    else:
                        yield from show_text(history, state, state['current_role'], message)

            except Exception as e:
                yield history, state, None, state['debate_state'][0], state['debate_state'][1], gr.update(interactive=False)

        # 新增人类输入处理函数
        def handle_human_input(human_text, history, state):
            if state.get('waiting_input'):
                history.append({"role": state['current_role'], "content": state['current_human_debater'] + ':\n' + human_text})
                state['current_human_debater_message'] = human_text
            yield history, state, None, state['debate_state'][0], state['debate_state'][1], ""

        def start_debate(state, play_voice, pro_role, mode, agent_type):
            state['running'] = True
            state['play_voice'] = play_voice
            state['pro_role'] = pro_role
            if agent_type == '简单智能体':
                state['agent_type'] = 'basic'
            elif agent_type == 'RAG反思智能体':
                state['agent_type'] = 'rag'
            else:
                state['agent_type'] = 'basic'
            # 人机辩论时，state记录辩论模式和人类属于正反方
            if mode == "人机辩论":
                state['debate_model'] = "ai_vs_human"
                state['human_size'] = "反方" if pro_role == "正方" else "正方"
            else:
                state['debate_model'] = "ai_vs_ai"
            return [
                [],  # 对应debate_screen
                state,
                gr.update(interactive=False),  # start_btn
                gr.update(interactive=True),  # stop_btn
                gr.update(interactive=False),  # debate_topic
                gr.update(interactive=False),  # debate_mode
                gr.update(interactive=False),  # agent_type
                gr.update(interactive=False),  # pro_role
                gr.update(interactive=False),  # con_role
                gr.update(interactive=False)

            ]
        def stop_debate(state):
            state['running'] = False
            return [

                state,
                gr.update(interactive=True),  # start_btn
                gr.update(interactive=False),  # stop_btn
                gr.update(interactive=True),  # debate_topic
                gr.update(interactive=True),  # debate_mode
                gr.update(interactive=True),  # agent_type
                gr.update(interactive=True),  # pro_role
                gr.update(interactive=True)  # con_role
            ]

        def transcribe(state, audio_chunk):
            if state['recording']:
                cache = state.get('record_cache', {})
                state['record_cache'] = splice_chunks(cache, audio_chunk)
            return state
        def stop_and_build(state,human_input):
            if state['recording']:
                cache = state.get('record_cache', {})
                if cache is not None and len(cache) > 0:
                    return state, human_input + "\n" + asr_decode(cache)
                else:
                    state['record_cache'] = np.array([])
            return state, human_input
        def start_record(state):
            state['recording'] = True
            return state
        # endregion

        # region 角色联动
        pro_role.change(
            update_role,
            pro_role,
            [con_role, pro_panel_label, right_panel_label]  # 添加两个标题输出
        )
        con_role.change(
            update_role,
            con_role,
            [pro_role, right_panel_label, pro_panel_label]  # 反向联动
        )

        # endregion

        # region 动态更新右侧面板
        debate_mode.change(
            update_right_panel,
            inputs=[debate_mode, debate_state],
            outputs=[human_interface, con_speakers,pro_role, con_role, debate_state]
        )

        agent_type.change(
            update_agent_list,
            agent_type,
            right_panel_label
        )

        is_play.change(
            lambda x: gr.update(visible=x),  # 使用 lambda 替代独立函数
            inputs=is_play,
            outputs=soundfile
        )

        # endregion

        # region 在控制按钮事件中更新状态
        start_btn.click(
            start_debate,
            inputs=[debate_state, is_play, pro_role, debate_mode, agent_type],
            outputs=[
                debate_screen,  # 清空聊天记录
                debate_state,  # 更新状态
                start_btn,  # 启用开始按钮
                stop_btn,  # 禁用停止按钮
                debate_topic,  # 清空辩论题目
                debate_mode,  # 启用辩论模式选择
                agent_type,  # 启用辩论智能体
                pro_role,  # 左侧角色选择
                con_role,
                is_play,
                # 右侧角色选择
            ],
            queue=False
        ).then(
            debate_generator,
            inputs=[debate_topic, debate_screen, debate_state],
            outputs=[debate_screen, debate_state, soundfile, pro_speakers, con_speakers,send_btn]
        )
        stop_btn.click(
            stop_debate,
            inputs=debate_state,
            outputs=[
                debate_state,
                start_btn,
                stop_btn,
                debate_topic,
                debate_mode,
                agent_type,
                pro_role,  # 左侧角色选择
                con_role]
        )
        reset_btn.click(
            reset_debate,
            inputs=debate_state,
            outputs=[
                debate_screen,  # 清空聊天记录
                debate_state,  # 更新状态
                start_btn,  # 启用开始按钮
                stop_btn,  # 禁用停止按钮
                debate_topic,  # 清空辩论题目
                debate_mode,  # 启用辩论模式选择
                agent_type,  # 启用辩论智能体
                pro_role,  # 左侧角色选择
                con_role,  # 右侧角色选择
                is_play
            ]
        )
        # 修改事件链
        # 修改事件链
        send_btn.click(
            handle_human_input,
            inputs=[human_input, debate_screen, debate_state],
            outputs=[debate_screen, debate_state, soundfile, pro_speakers, con_speakers,human_input]
        ).then(
            debate_generator,  # 直接重新调用生成器函数
            inputs=[debate_topic, debate_screen, debate_state],
            outputs=[debate_screen, debate_state, soundfile, pro_speakers, con_speakers,send_btn]
        )

        record_audio.stream(
            transcribe,
            [debate_state, record_audio],
            debate_state,
            show_progress='minimal',
            trigger_mode="always_last"  # 确保最后触发一次处理
        )

        record_audio.stop_recording(stop_and_build, [debate_state,human_input], [debate_state, human_input])
        record_audio.start_recording(start_record, debate_state, debate_state)
    # endregion

    demo.launch()


if __name__ == "__main__":
    main()
